from RumdetecFramework.DistillRumorFramework import SentDistillModel
from Dataloader.twitterloader import *
from Dataloader.dataloader_utils import *
from SentModel.Sent2Vec import *
from PropModel.SeqPropagation import GRUModel

def obtain_DistillBertRD(pretrained_file):
    ch = torch.load(pretrained_file)
    bvec = BertVec("../../bert_en/", bert_parallel=True)
    lvec = LSTMVec("../saved/bert_embedding.pkl", "../../bert_en/", 768, 768)
    prop = GRUModel(768, 256, 1, 0.2)
    cls = nn.Linear(256, 2)
    bvec.bert.load_state_dict(ch['bert'])
    prop.load_state_dict(ch['rmdModel'])
    cls.load_state_dict(ch['rdm_classifier'])
    SDRD = SentDistillModel(lvec, prop, cls, bvec, distill_level='token')
    return SDRD

def obtain_general_set(tr_prefix, dev_prefix, te_prefix):
    tr_set = TwitterSet()
    tr_set.load_data_fast(data_prefix=tr_prefix, min_len=5)
    dev_set = TwitterSet()
    dev_set.load_data_fast(data_prefix=dev_prefix)
    te_set = TwitterSet()
    te_set.load_data_fast(data_prefix=te_prefix)
    return tr_set, dev_set, te_set

# for i in range(5):
i = 2
tr, dev, te = obtain_general_set("../data/twitter_tr%d"%i, "../data/twitter_dev%d"%i, "../data/twitter_te%d"%i)
tr.filter_short_seq(min_len=5)
tr.trim_long_seq(10)
print("%s : (dev event)/(test event)/(train event) = %3d/%3d/%3d" % (te.data[te.data_ID[0]]['event'], len(dev), len(te), len(tr)))
print("\n\n===========DistillBertRD Train===========\n\n")
model = obtain_DistillBertRD("../saved/BertRDM_best.pkl")
model.train_iters(tr, dev, te, max_epochs=100,
                log_dir="../logs/", log_suffix="DistillBertRD_%s"%(te.data[te.data_ID[0]]['event']),
                model_file="DistillBertRD_%s.pkl"% (te.data[te.data_ID[0]]['event']))